{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "dd4c8776-6853-4257-bef8-72778724ad57",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "# Embedding Text with Azure OpenAI\n",
    "\n",
    "The Azure OpenAI service can be used to solve a large number of natural language tasks through prompting the completion API. To make it easier to scale your prompting workflows from a few examples to large datasets of examples we have integrated the Azure OpenAI service with the distributed machine learning library [SynapseML](https://www.microsoft.com/en-us/research/blog/synapseml-a-simple-multilingual-and-massively-parallel-machine-learning-library/). This integration makes it easy to use the [Apache Spark](https://spark.apache.org/) distributed computing framework to process millions of prompts with the OpenAI service. This tutorial shows how to apply large language models to generate embeddings for large datasets of text. \n",
    "\n",
    "## Step 1: Prerequisites\n",
    "\n",
    "The key prerequisites for this quickstart include a working Azure OpenAI resource, and an Apache Spark cluster with SynapseML installed. We suggest creating a Synapse workspace, but an Azure Databricks, HDInsight, or Spark on Kubernetes, or even a python environment with the `pyspark` package will work. \n",
    "\n",
    "1. An Azure OpenAI resource – request access [here](https://customervoice.microsoft.com/Pages/ResponsePage.aspx?id=v4j5cvGGr0GRqy180BHbR7en2Ais5pxKtso_Pz4b1_xUOFA5Qk1UWDRBMjg0WFhPMkIzTzhKQ1dWNyQlQCN0PWcu) before [creating a resource](https://docs.microsoft.com/en-us/azure/cognitive-services/openai/how-to/create-resource?pivots=web-portal#create-a-resource)\n",
    "1. [Create a Synapse workspace](https://docs.microsoft.com/en-us/azure/synapse-analytics/get-started-create-workspace)\n",
    "1. [Create a serverless Apache Spark pool](https://docs.microsoft.com/en-us/azure/synapse-analytics/get-started-analyze-spark#create-a-serverless-apache-spark-pool)\n",
    "\n",
    "\n",
    "## Step 2: Import this guide as a notebook\n",
    "\n",
    "The next step is to add this code into your Spark cluster. You can either create a notebook in your Spark platform and copy the code into this notebook to run the demo. Or download the notebook and import it into Synapse Analytics\n",
    "\n",
    "1.\t[Download this demo as a notebook](https://github.com/microsoft/SynapseML/blob/master/notebooks/features/cognitive_services/CognitiveServices%20-%20OpenAI%20Embedding.ipynb) (click Raw, then save the file)\n",
    "1.\tImport the notebook [into the Synapse Workspace](https://docs.microsoft.com/en-us/azure/synapse-analytics/spark/apache-spark-development-using-notebooks#create-a-notebook) or if using Databricks [into the Databricks Workspace](https://docs.microsoft.com/en-us/azure/databricks/notebooks/notebooks-manage#create-a-notebook)\n",
    "1. Install SynapseML on your cluster. Please see the installation instructions for Synapse at the bottom of [the SynapseML website](https://microsoft.github.io/SynapseML/). Note that this requires pasting an additional cell at the top of the notebook you just imported\n",
    "3.\tConnect your notebook to a cluster and follow along, editing and rnnung the cells below.\n",
    "\n",
    "## Step 3: Fill in your service information\n",
    "\n",
    "Next, please edit the cell in the notebook to point to your service. In particular set the `service_name`, `deployment_name`, `location`, and `key` variables to match those for your OpenAI service:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "1b0db8af-7fe2-40bc-9df4-cc7f274d53f0",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "from synapse.ml.core.platform import find_secret\n",
    "\n",
    "# Fill in the following lines with your service information\n",
    "# Learn more about selecting which embedding model to choose: https://openai.com/blog/new-and-improved-embedding-model\n",
    "service_name = \"synapseml-openai-2\"\n",
    "deployment_name_embeddings = \"text-embedding-ada-002\"\n",
    "\n",
    "key = find_secret(\n",
    "    secret_name=\"openai-api-key-2\", keyvault=\"mmlspark-build-keys\"\n",
    ")  # please replace this with your key as a string\n",
    "\n",
    "assert key is not None and service_name is not None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "22603456-4f44-4b3a-9751-9ca5231b799b",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "## Step 4: Load Data\n",
    "\n",
    "In this demo we will explore a dataset of fine food reviews"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "e93c3f26-95c5-4708-8630-b71620841fc5",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "import pyspark.sql.functions as F\n",
    "\n",
    "df = (\n",
    "    spark.read.options(inferSchema=\"True\", delimiter=\",\", header=True)\n",
    "    .csv(\"wasbs://publicwasb@mmlspark.blob.core.windows.net/fine_food_reviews_1k.csv\")\n",
    "    .repartition(5)\n",
    ")\n",
    "\n",
    "df = df.withColumn(\n",
    "    \"combined\",\n",
    "    F.format_string(\"Title: %s; Content: %s\", F.trim(df.Summary), F.trim(df.Text)),\n",
    ")\n",
    "\n",
    "display(df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "0c69ee56-172f-413b-a335-d15482fda55e",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "## Step 5: Generate Embeddings\n",
    "\n",
    "We will first generate embeddings for the reviews using the SynapseML OpenAIEmbedding client."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "0cfffbb5-2ca0-4e58-b86d-0c89195f7f5a",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "from synapse.ml.services.openai import OpenAIEmbedding\n",
    "\n",
    "embedding = (\n",
    "    OpenAIEmbedding()\n",
    "    .setSubscriptionKey(key)\n",
    "    .setDeploymentName(deployment_name_embeddings)\n",
    "    .setCustomServiceName(service_name)\n",
    "    .setTextCol(\"combined\")\n",
    "    .setErrorCol(\"error\")\n",
    "    .setOutputCol(\"embeddings\")\n",
    ")\n",
    "\n",
    "completed_df = embedding.transform(df).cache()\n",
    "display(completed_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "7a46d8be-3f0a-43dc-b594-5542fc6c3620",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "## Step 6: Reduce Embedding dimensionality for Visualization\n",
    "We reduce the dimensionality to 2 dimensions using t-SNE decomposition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "f5358431-90c5-49d1-a4c7-b411c435a664",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from sklearn.manifold import TSNE\n",
    "import numpy as np\n",
    "\n",
    "collected = list(completed_df.collect())\n",
    "matrix = np.array([[r[\"embeddings\"]] for r in collected])[:, 0, :].astype(np.float64)\n",
    "scores = np.array([[r[\"Score\"]] for r in collected]).reshape(-1)\n",
    "\n",
    "tsne = TSNE(n_components=2, perplexity=15, random_state=42, init=\"pca\")\n",
    "vis_dims = tsne.fit_transform(matrix)\n",
    "vis_dims.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "669cba60-c36b-409d-b8c0-5da04c676091",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "## Step 7: Plot the embeddings\n",
    "\n",
    "We now use t-SNE to reduce the dimensionality of the embeddings from 1536 to 2. Once the embeddings are reduced to two dimensions, we can plot them in a 2D scatter plot. We colour each review by its star rating, ranging from red for negative reviews, to green for positive reviews. We can observe a decent data separation even in the reduced 2 dimensions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "9240e325-7db1-43fc-9340-c7b1ba303adb",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import matplotlib\n",
    "import numpy as np\n",
    "\n",
    "colors = [\"red\", \"darkorange\", \"gold\", \"turquoise\", \"darkgreen\"]\n",
    "x = [x for x, y in vis_dims]\n",
    "y = [y for x, y in vis_dims]\n",
    "color_indices = scores - 1\n",
    "\n",
    "colormap = matplotlib.colors.ListedColormap(colors)\n",
    "plt.scatter(x, y, c=color_indices, cmap=colormap, alpha=0.3)\n",
    "for score in [0, 1, 2, 3, 4]:\n",
    "    avg_x = np.array(x)[scores - 1 == score].mean()\n",
    "    avg_y = np.array(y)[scores - 1 == score].mean()\n",
    "    color = colors[score]\n",
    "    plt.scatter(avg_x, avg_y, marker=\"x\", color=color, s=100)\n",
    "\n",
    "plt.title(\"Amazon ratings visualized in language using t-SNE\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "bc70f837-b473-4918-aaf0-10db4c255896",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "## Step 8: Build a fast vector index to over review embeddings\n",
    "\n",
    "We will use SynapseML's KNN estimator to build a fast cosine-similarity retrieval engine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "30ac70a5-40b9-4818-9e32-69512407c8c1",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "from synapse.ml.nn import *\n",
    "\n",
    "knn = (\n",
    "    KNN()\n",
    "    .setFeaturesCol(\"embeddings\")\n",
    "    .setValuesCol(\"id\")\n",
    "    .setOutputCol(\"output\")\n",
    "    .setK(10)\n",
    ")\n",
    "\n",
    "knn_index = knn.fit(completed_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "6885033f-6eea-4338-a632-2837582d91a1",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "## Step 8: Build the retrieval model pipeline\n",
    "\n",
    "Note: The data types of the ID columns in the document and query dataframes should be the same. For some OpenAI models, users should use separate models for embedding documents and queries. These models are denoted by the \"-doc\" and \"-query\" suffixes respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "72558101-d530-43ed-adfc-53e3f332d7d9",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "from pyspark.ml import PipelineModel\n",
    "\n",
    "embedding_query = (\n",
    "    OpenAIEmbedding()\n",
    "    .setSubscriptionKey(key)\n",
    "    .setDeploymentName(deployment_name_embeddings)\n",
    "    .setCustomServiceName(service_name)\n",
    "    .setTextCol(\"query\")\n",
    "    .setErrorCol(\"error\")\n",
    "    .setOutputCol(\"embeddings\")\n",
    ")\n",
    "\n",
    "retrieval_model = PipelineModel(stages=[embedding_query, knn_index])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "32c7b964-0c68-4ab6-b656-865cb59ef577",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "## Step 9: Retrieve results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "c3054be0-9b1a-49ae-95da-309751f2ad0a",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "query_df = (\n",
    "    spark.createDataFrame(\n",
    "        [\n",
    "            (\n",
    "                0,\n",
    "                \"desserts\",\n",
    "            ),\n",
    "            (\n",
    "                1,\n",
    "                \"disgusting\",\n",
    "            ),\n",
    "        ]\n",
    "    )\n",
    "    .toDF(\"id\", \"query\")\n",
    "    .withColumn(\"id\", F.col(\"id\").cast(\"int\"))\n",
    ")\n",
    "\n",
    "\n",
    "df_matches = retrieval_model.transform(query_df).cache()\n",
    "\n",
    "df_result = (\n",
    "    df_matches.withColumn(\"match\", F.explode(\"output\"))\n",
    "    .join(df, df[\"id\"] == F.col(\"match.value\"))\n",
    "    .select(\"query\", F.col(\"combined\"), \"match.distance\")\n",
    ")\n",
    "\n",
    "display(df_result)"
   ]
  }
 ],
 "metadata": {
  "application/vnd.databricks.v1+notebook": {
   "dashboards": [],
   "language": "python",
   "notebookMetadata": {
    "pythonIndentUnit": 2
   },
   "notebookName": "AI Services - OpenAI Embedding",
   "notebookOrigID": 2846029038781858,
   "widgets": {}
  },
  "kernel_info": {
   "name": "synapse_pyspark"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  },
  "save_output": true,
  "synapse_widget": {
   "state": {
    "4bd0e60b-98ae-4bfe-98ee-6f0399ceb456": {
     "persist_state": {
      "view": {
       "chartOptions": {
        "aggregationType": "count",
        "categoryFieldKeys": [
         "0"
        ],
        "chartType": "bar",
        "isStacked": false,
        "seriesFieldKeys": [
         "0"
        ]
       },
       "tableOptions": {},
       "type": "details"
      }
     },
     "sync_state": {
      "isSummary": false,
      "language": "scala",
      "table": {
       "rows": [
        {
         "0": "Once upon a time",
         "1": [
          " there was a girl who had a dream of becoming a writer.\n\nShe started writing short stories"
         ]
        },
        {
         "0": "Hello my name is",
         "1": [
          "***** and I have a question about my cat\n\nHello, thank you for bringing your question to"
         ]
        },
        {
         "0": "The best code is code thats",
         "1": [
          " not there\n\nCommenting your code is important. Not only does it help you remember what you"
         ]
        }
       ],
       "schema": [
        {
         "key": "0",
         "name": "prompt",
         "type": "string"
        },
        {
         "key": "1",
         "name": "text",
         "type": "ArrayType(StringType,true)"
        }
       ],
       "truncated": false
      }
     },
     "type": "Synapse.DataFrame"
    }
   },
   "version": "0.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
